import os  # 操作系统接口
import pandas as pd  # 数据处理和分析
import numpy as np  # 数值计算
import re  # 正则表达式
from typing import Dict, List, Tuple, Optional  # 类型提示
from datetime import datetime  # 日期时间处理


def extract_class_from_filename(filename: str) -> str:
    """
    Extract class label from filename.
    
    Examples:
    - angle-scl3300d01-empty-0.csv -> empty
    - angle-scl3300d01-sitting-0.csv -> sitting
    - angle-scl3300d01-sitting-laydown-0.csv -> sitting-laydown
    """
    # Extract the part between 'angle-scl3300d01-' and the last dash
    match = re.search(r'angle-scl3300d01-(.*)-\d+', filename)  # 使用正则表达式匹配文件名模式
    if match:  # 如果匹配成功
        return match.group(1)  # 返回捕获的类名部分
    return "unknown"  # 默认返回"unknown"


def load_data(data_dir: str) -> Tuple[Dict[str, pd.DataFrame], List[str]]:
    """
    Load all CSV files from the data directory.
    
    Args:
        data_dir: Path to directory containing CSV files
        
    Returns:
        Dictionary mapping class names to DataFrames and list of unique classes
    """
    # Convert relative path to absolute path if needed
    if not os.path.isabs(data_dir):  # 检查是否为绝对路径
        script_dir = os.path.dirname(os.path.abspath(__file__))  # 获取当前脚本目录
        data_dir = os.path.abspath(os.path.join(script_dir, data_dir))  # 转换为绝对路径
    
    if not os.path.exists(data_dir):  # 检查数据目录是否存在
        raise FileNotFoundError(f"Data directory not found: {data_dir}")  # 抛出异常
    
    data_by_class = {}  # 按类别存储数据的字典
    classes = []  # 存储所有类别的列表
    
    for filename in os.listdir(data_dir):  # 遍历数据目录
        if filename.endswith('.csv') and not filename.startswith('.'):  # 只处理CSV文件
            file_path = os.path.join(data_dir, filename)  # 构建完整文件路径
            class_name = extract_class_from_filename(filename)  # 从文件名提取类别
            
            if class_name not in classes:  # 如果是新类别
                classes.append(class_name)  # 添加到类别列表
            
            # Read the CSV file
            df = pd.read_csv(file_path)  # 读取CSV文件
            
            # Convert timestamp to datetime
            if 'ts' in df.columns:  # 如果有时间戳列
                df['timestamp'] = pd.to_datetime(df['ts'], unit='ms')  # 转换为datetime
                df = df.sort_values('timestamp')  # 按时间排序
            
            # Add the class label to the DataFrame
            df['class'] = class_name  # 添加类别标签列
            
            if class_name in data_by_class:  # 如果已有该类别数据
                data_by_class[class_name] = pd.concat([data_by_class[class_name], df])  # 合并数据
            else:
                data_by_class[class_name] = df  # 创建新类别条目
    
    return data_by_class, classes


def segment_time_series(df: pd.DataFrame, window_size: int, step_size: int) -> List[pd.DataFrame]:
    """
    Segment time series data into windows.
    
    Args:
        df: DataFrame containing time series data
        window_size: Size of the window in number of samples
        step_size: Step size for sliding window
        
    Returns:
        List of DataFrames, each containing a window of data
    """
    segments = []  # 存储所有数据段的列表
    for i in range(0, len(df) - window_size + 1, step_size):  # 使用滑动窗口遍历数据
        segment = df.iloc[i:i + window_size].copy()  # 获取当前窗口的数据副本
        segments.append(segment)  # 将数据段添加到结果列表
    
    return segments


def prepare_dataset(data_dir: str, window_size: int = 100, step_size: int = 50) -> Tuple[List[pd.DataFrame], List[str], List[str]]:
    """
    Prepare dataset by loading data, segmenting it, and creating labels.
    
    Args:
        data_dir: Path to directory containing CSV files
        window_size: Size of the window in number of samples
        step_size: Step size for sliding window
        
    Returns:
        List of DataFrames (segments), list of corresponding labels, and list of unique classes
    """
    data_by_class, classes = load_data(data_dir)  # 加载原始数据并按类别分组
    
    segments = []  # 存储所有数据段的列表
    labels = []  # 存储对应标签的列表
    
    for class_name, df in data_by_class.items():  # 遍历每个类别的数据
        # Select only the sensor data columns
        sensor_cols = ['acc_x', 'acc_y', 'acc_z', 'angle_x', 'angle_y', 'angle_z']  # 传感器数据列名
        df_sensors = df[sensor_cols]  # 只保留传感器数据列
        
        # Segment the time series
        class_segments = segment_time_series(df_sensors, window_size, step_size)  # 对当前类别数据进行分段
        
        segments.extend(class_segments)  # 将数据段添加到总列表
        labels.extend([class_name] * len(class_segments))  # 为每个数据段添加对应标签
    
    return segments, labels, classes


if __name__ == "__main__":
    # Test the data loader
    data_dir = "../data"
    segments, labels, classes = prepare_dataset(data_dir)
    print(f"Loaded {len(segments)} segments with {len(classes)} classes: {classes}")
    print(f"First segment shape: {segments[0].shape}")
